[None][Perf] Resolve the runtime JIT issue - #18813

Three TunableRunners exist only to drive DeepGemm's JIT as a side effect of
autotuning: fp8SwapABGemmRunner, Fp8PrequantizedSwapABGemmRunner and
Fp8BlockScalingGemmRunner. All three are single-tactic, so each bucket costs
one forward call and nothing else.

They warm up over deep_gemm_gen_tuning_buckets, a grid built for tactic
selection: step 8 below 128, step 128 above. DeepGemm picks its tile layout
from a heuristic over M and generates one kernel per layout, so where that
grid never samples an M that selects a layout, the layout is compiled
mid-inference instead. DeepGemm forks nvcc while holding the GIL, so one
rank's compile stalls every rank of its attention-DP group.

The gap is transition width in the low-M region. Per get_layout_candidates in
deepgemm/csrc/jit_kernels/heuristics/sm100.hpp, the swap-AB branch enumerates
block_m in multiples of lcm(16, block_m_multiple_of) -- steps of 16 -- up to
256, with block_n pinned to 128; the non-swap branch switches block_m at
m <= 32 and m <= 64. So below ~2K the selected layout moves roughly every 16
tokens of M, where the old grid steps 128. Above ~2K block_m has saturated
and the selection moves only at coarse granularity, so step 128 suffices.

Give the three runners their own grid matched to that: step 8 below 128,
step 16 to 2048, step 128 above, plus the bound itself. Below 128 the
grid stops at the low band, as the old one does, so a worker whose M never
leaves that band is unaffected -- see the GEN worker below.

Measured on B200 (SM100, 148 SMs), one process and one fresh JIT cache per
arm: warm up over the grid, then sweep every M in 1..8192 and count kernel
directories created during the sweep. Each such directory is one nvcc compile
that would have landed under traffic.

  grid            buckets   7168x2048  2048x7168  576x7168  4608x7168
  step 8/128           78           5          3         1          3
  this grid           184           0          0         0          0

Zero residual compiles on all four shapes, against 12 across the four for the
old grid. The residuals all land just above a multiple of 16 -- 129, 145, 161,
193, 257, 321, 385, 513, 769 -- which is the 16-step block_m enumeration
showing through, and is what sampling at 16 covers.

Independently, through the production path, with a wall clock on it:
torch.ops.trtllm.fp8_swap_ab_gemm -- the op FP8BlockScalesLinearMethod.apply
dispatches to on SM100 -- warmed by the real AutoTuner over the runner's own
tuning_config, then every M in 1..8192 issued and individually timed. A call
that takes seconds rather than microseconds is nvcc running inline in the GEMM,
and each is confirmed by a JIT cache-directory delta.

  shape       old warmup   old stalls  new warmup   new stalls
  7168x2048       23.4 s    5 / 10.06 s    35.0 s    0 / 0 s
  2048x7168       38.5 s    3 /  6.05 s    44.7 s    0 / 0 s
  4608x7168       29.1 s    3 /  5.46 s    34.4 s    0 / 0 s
  576x7168        60.9 s    1 /  1.85 s    57.2 s    0 / 0 s
  total                    12 / 23.41 s              0 / 0 s

Same 12 compiles the sweep counted, at the same M -- 129, 161, 193, 257, 321;
257, 513, 769; 145, 257, 385; 513 -- now with a cost: 1.8-2.3 s each, 23.41 s
in total, spent inside a GEMM call.

The grid is denser than the one it replaces, 184 buckets against 78, and the
cost of that is the compiles it hoists rather than the buckets themselves.
Warmup grows 19.4 s across the four shapes to remove 23.41 s of in-GEMM stall,
i.e. the compiles are not made cheaper, they are moved off the serving path.
That is the whole benefit: in place they run while DeepGemm holds the GIL, so
every rank of the attention-DP group waits on them, whereas at startup they are
paid once alongside the rest of engine build. The ~106 extra buckets that land
on an already-compiled layout are a GEMM launch each and cost little enough
that on 576x7168 the denser grid warmed up 3.7 s faster than the sparser one --
run-to-run nvcc variance exceeds them. A single-tactic runner takes the
autotuner's single-pair path, so no bucket runs the timed profile loop. The
runtime GEMM path is untouched: the heuristic still sees the true M, so no tile
choice changes.

End to end on a real engine, DeepSeek-V3-Lite/fp8 (block-scaled shapes
12288x2560 and 2560x12288), fresh JIT cache: the new grid compiles 9 more
sm100_fp8_fp4_gemm_1d1d kernels during engine build than the old one, 108
against 99, with the non-GEMM kernel count unchanged. Those 9 are layouts the
old grid never samples, so under it they land under traffic instead -- about
17 s of GIL-held stall at the ~1.9 s per compile measured above, on a model far
smaller than K3.

And on the deployment this was written for, Kimi-K3-NVFP4 disaggregated on
GB300 -- a 16-rank attention-DP context worker, four clean runs without the fix
and five with it. DG_PRINT_CONFIGS says the context phase issues eight fp8 GEMM
shapes, over which M up to 8192 selects 142 distinct layouts; the old grid
primes 116, and every one of the 26 it misses is selected by an M between 129
and 1760, i.e. by a partial chunk. With DG_JIT_PRINT_COMPILER_COMMAND counting
forks, three instrumented arms without the fix ran nvcc 59, 63 and 49 times
after serving began, against 0, 0 and 0 with it, across two cluster
configurations -- and every count was predicted from that run's own token stream
before its log was read. Context steps that forked nvcc took 3.66 s at the
median against 1.01 s for steps that did not.

The same trade shows up there as on the microbenchmarks, and the cost side is
worth stating precisely because it is easy to quote the wrong worker's number.
Summing rank 0's [Autotuner] Autotuning process start/end intervals, four runs
without the change against six with it. Every CTX range is non-overlapping
(449-463 against 513-531); the GEN column's ON side is the single run built on
the final guarded grid:

  worker   autotune       process startup
  CTX      454 -> 523 s   1106 -> 1173 s
  GEN      126 -> 129 s    900 ->  914 s

Server-ready is max(CTX, GEN), and with GEN at 914 s the context worker gates,
so the end-to-end cost is CTX's: 1106 s -> 1173 s, +67 s. Roughly one more
minute of engine build, once, in exchange for every compile that would otherwise
have landed mid-inference with the GIL held. Over the same runs the denser grid
forks nvcc 2301 times against 2036.

GEN barely moves because of the sub-128 guard, and that is the whole reason the
guard is there. The bound passed in is tune_max_num_tokens where the runner sets
one, and the warmup forward's actual M otherwise, which for a disagg generation
worker is batch x MTP -- 32 on this deployment, against 8192 for the context
worker. Clamping that up to the 4096 floor unconditionally, as an earlier
revision of this patch did, gives GEN 152 buckets of which it can reach four;
measured on the same config that cost 409 s of autotune and 1163 s of startup,
with [Autotuner] Cache size after warmup at 1342 against 246. With the guard GEN
gets the same 15 buckets the old function gave it, and its cache size is 246 --
bit-identical to the old grid, while CTX's is 1672 either way, so the guard
provably moves only the worker it is meant to. That matters beyond the wasted
time: GEN otherwise finishes ~200 s ahead of CTX, and spending that slack would
have put GEN on the critical path -- it already gated one of the unguarded runs,
GEN 1162 s against CTX 1148 s.

End to end across those nine runs TTFT moves where that mechanism says it
should and nowhere else. The low percentiles separate cleanly, with no run
overlapping the other arm's range: p5 5.50 s -> 4.07 s, -26% (5.11-6.39 against
3.99-4.27), and p10 8.64 s -> 5.42 s, -37% (8.01-9.66 against 5.15-5.64). The
mean and the output token throughput also separate, but only just, by 28 ms and
by 5 tok/s of margin between the arms' ranges: 63.5 s -> 62.1 s (-2.3%) and
1677 -> 1701 tok/s (+1.4%, 1668-1691 against 1696-1708). The median does not
move at all, 65.9 s against 65.6 s, p25 and p90 stay inside their own spread,
and p99 changes sign between samples, so treat it as noise. That shape is what
the mechanism predicts rather than a weakness in it: a compile costs ~2.6 s and
lands in one step, so it dominates the TTFT of a request waiting on a single
context step and vanishes into the queueing of one waiting on hundreds.

Worth recording what this does not show, because it is the obvious thing to
claim and it does not hold: of the 11 slowest context steps on the arms without
the fix only 2 contain a compile, so the drop in steps over 10 s -- 5 and 6
against 1 and 2 on one configuration, 11 against 0 on the other -- is larger
than the compiles can account for and is partly run-to-run variance.

Fp8BlockScalingGemmRunner was also missing exclude_from_cache, which the other
two already set. Without it a warm disk cache makes choose_one() return the
stored tactic and skip the tuning path outright, so its JIT warmup never ran
on any startup after the first -- the compiles it exists to hoist landed under
traffic anyway. It is single-tactic and its forward() ignores `tactic`, so
nothing is lost by not persisting it. This is an unrelated bug fix rather than
part of either cause above, and it is carried here because the two share every
line of context. Note this runner's warmup is Hopper-only in effect: on SM100
the block-scaling GEMM dispatches to TrtllmGenGemmRunner's prebuilt cubins and
compiles nothing.

deep_gemm_gen_tuning_buckets is left as-is for its two CuTe DSL consumers,
which are a different compiler with a separate cache, and which tune under the
PARALLEL strategy where a denser grid would go through the full timed profile
loop rather than the single-pair path.

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>